#
# This source file is part of the EdgeDB open source project.
#
# Copyright 2020-present MagicStack Inc. and the EdgeDB authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from __future__ import annotations
from typing import Any, Callable

import collections
import functools
import json
import uuid

import immutables

from edb.common import checked
from edb.common import verutils
from edb.common import uuidgen

from edb.schema import abc as s_abc
from edb.schema import expr as s_expr
from edb.schema import functions as s_func
from edb.schema import name as s_name
from edb.schema import objects as s_obj
from edb.schema import operators as s_oper
from edb.schema import schema as s_schema
from edb.schema import version as s_ver

from . import structure as sr_struct


SchemaClassLayout = dict[type[s_obj.Object], sr_struct.SchemaTypeLayout]


def parse_schema(
    base_schema: s_schema.Schema,
    data: str | bytes,
    schema_class_layout: SchemaClassLayout,
) -> s_schema.FlatSchema:
    """Parse JSON-encoded schema objects and populate the schema with them.

    Args:
        schema:
            A schema instance to use as a starting point.
        data:
            A JSON-encoded schema object data as returned
            by an introspection query.
        schema_class_layout:
            A mapping describing schema class layout in the reflection,
            as returned from
            :func:`schema.reflection.structure.generate_structure`.

    Returns:
        A schema instance including objects encoded in the provided
        JSON sequence.
    """

    id_to_type = {}
    id_to_data = {}
    name_to_id = {}
    shortname_to_id = collections.defaultdict(set)
    globalname_to_id = {}
    dict_of_dicts: Callable[
        [],
        dict[tuple[type[s_obj.Object], str], dict[uuid.UUID, None]],
    ] = functools.partial(collections.defaultdict, dict)
    refs_to: dict[
        uuid.UUID,
        dict[tuple[type[s_obj.Object], str], dict[uuid.UUID, None]]
    ] = collections.defaultdict(dict_of_dicts)

    objects: dict[uuid.UUID, tuple[s_obj.Object, dict[str, Any]]] = {}
    objid: uuid.UUID

    for entry in json.loads(data):
        _, _, clsname = entry['_tname'].rpartition('::')
        mcls = s_obj.ObjectMeta.maybe_get_schema_class(clsname)
        if mcls is None:
            raise ValueError(
                f'unexpected type in schema reflection: {clsname}')
        objid = uuidgen.UUID(entry['id'])
        objects[objid] = (mcls._create_from_id(objid), entry)

    refdict_updates = {}

    for objid, (obj, entry) in objects.items():
        mcls = type(obj)
        name = s_name.name_from_string(entry['name__internal'])
        layout = schema_class_layout[mcls]

        if (
            base_schema.has_object(objid)
            and not isinstance(obj, s_ver.BaseSchemaVersion)
        ):
            continue

        if isinstance(obj, s_obj.QualifiedObject):
            name_to_id[name] = objid
        else:
            name = s_name.UnqualName(str(name))
            globalname_to_id[mcls, name] = objid

        if isinstance(obj, (s_func.Function, s_oper.Operator)):
            shortname = mcls.get_shortname_static(name)
            shortname_to_id[mcls, shortname].add(objid)

        id_to_type[objid] = type(obj).__name__

        all_fields = mcls.get_schema_fields()
        objdata: list[Any] = [None] * len(all_fields)
        val: Any
        refid: uuid.UUID

        for k, v in entry.items():
            desc = layout.get(k)
            if desc is None:
                continue

            fn = desc.fieldname
            field = all_fields.get(fn)
            if field is None:
                continue
            findex = field.index

            if desc.storage is not None:
                if v is None:
                    pass
                elif desc.storage.ptrkind == 'link':
                    refid = uuidgen.UUID(v['id'])
                    newobj = objects.get(refid)
                    if newobj is not None:
                        val = newobj[0]
                    else:
                        val = base_schema.get_by_id(refid)
                    objdata[findex] = val.schema_reduce()
                    refs_to[val.id][mcls, fn][objid] = None

                elif desc.storage.ptrkind == 'multi link':
                    ftype = mcls.get_field(fn).type
                    if issubclass(ftype, s_obj.ObjectDict):
                        refids = ftype._container(
                            uuidgen.UUID(e['value']) for e in v)
                        refkeys = tuple(e['name'] for e in v)
                        val = ftype(refids, refkeys, _private_init=True)
                    else:
                        refids = ftype._container(
                            uuidgen.UUID(e['id']) for e in v)
                        val = ftype(refids, _private_init=True)
                    objdata[findex] = val.schema_reduce()
                    for refid in refids:
                        refs_to[refid][mcls, fn][objid] = None

                elif desc.storage.shadow_ptrkind:
                    val = entry[f'{k}__internal']
                    ftype = mcls.get_field(fn).type
                    if val is not None and type(val) is not ftype:
                        if issubclass(ftype, s_expr.Expression):
                            val = _parse_expression(val, objid, k)
                            for refid in val.refs.ids():
                                refs_to[refid][mcls, fn][objid] = None
                        elif issubclass(ftype, s_expr.ExpressionList):
                            exprs = []
                            for e_dict in val:
                                e = _parse_expression(e_dict, objid, k)
                                assert e.refs is not None
                                for refid in e.refs.ids():
                                    refs_to[refid][mcls, fn][objid] = None
                                exprs.append(e)
                            val = ftype(exprs)
                        elif issubclass(ftype, s_expr.ExpressionDict):
                            expr_dict = dict()
                            for e_dict in val:
                                e = _parse_expression(
                                    e_dict['expr'], objid, k)
                                assert e.refs is not None
                                for refid in e.refs.ids():
                                    refs_to[refid][mcls, fn][objid] = None
                                expr_dict[e_dict['name']] = e
                            val = ftype(expr_dict)
                        elif issubclass(ftype, s_obj.Object):
                            val = val.id
                        elif issubclass(ftype, s_name.Name):
                            if isinstance(obj, s_obj.QualifiedObject):
                                val = s_name.name_from_string(val)
                            else:
                                val = s_name.UnqualName(val)
                        else:
                            val = ftype(val)

                    if issubclass(ftype, s_abc.Reducible):
                        val = val.schema_reduce()
                    objdata[findex] = val

                else:
                    ftype = mcls.get_field(fn).type
                    if type(v) is not ftype:
                        if issubclass(ftype, verutils.Version):
                            objdata[findex] = _parse_version(v)
                        elif issubclass(ftype, s_name.Name):
                            objdata[findex] = s_name.name_from_string(v)
                        elif (
                            issubclass(ftype, checked.ParametricContainer)
                            and ftype.types
                            and len(ftype.types) == 1
                        ):
                            # Coerce the elements in a parametric container
                            # type.
                            # XXX: Or should we do it in the container?
                            subtyp = ftype.types[0]
                            objdata[findex] = ftype(
                                subtyp(x) for x in v)  # type: ignore
                        else:
                            objdata[findex] = ftype(v)
                    else:
                        objdata[findex] = v

            elif desc.is_refdict:
                ftype = mcls.get_field(fn).type
                refids = ftype._container(uuidgen.UUID(e['id']) for e in v)
                for refid in refids:
                    refs_to[refid][mcls, fn][objid] = None

                val = ftype(refids, _private_init=True)
                objdata[findex] = val.schema_reduce()
                if desc.properties:
                    for e_dict in v:
                        refdict_updates[uuidgen.UUID(e_dict['id'])] = {
                            p: pv for p in desc.properties
                            if (pv := e_dict[f'@{p}']) is not None
                        }

        id_to_data[objid] = tuple(objdata)

    for objid, updates in refdict_updates.items():
        if updates:
            sclass = s_obj.ObjectMeta.get_schema_class(id_to_type[objid])
            updated_data = list(id_to_data[objid])
            for fn, v in updates.items():
                field = sclass.get_schema_field(fn)
                updated_data[field.index] = v
            id_to_data[objid] = tuple(updated_data)

    refs_to_im = {}
    for referred_id, ref_data in refs_to.items():
        refs_to_im[referred_id] = immutables.Map((
            (k, immutables.Map(r)) for k, r in ref_data.items()
        ))

    return s_schema.FlatSchema()._replace(
        id_to_type=immutables.Map(id_to_type),
        id_to_data=immutables.Map(id_to_data),
        name_to_id=immutables.Map(name_to_id),
        shortname_to_id=immutables.Map({
            (k, frozenset(v)) for k, v in shortname_to_id.items()
        }),
        globalname_to_id=immutables.Map(globalname_to_id),
        refs_to=immutables.Map(refs_to_im),
    )


def _parse_expression(
    val: dict[str, Any], id: uuid.UUID, field: str
) -> s_expr.Expression:
    refids = frozenset(
        uuidgen.UUID(r) for r in val['refs']
    )
    expr = s_expr.Expression(
        text=val['text'],
        refs=s_obj.ObjectSet(
            refids,
            _private_init=True,
        ),
    )
    expr.set_origin(id, field)
    return expr


def _parse_version(val: dict[str, Any]) -> verutils.Version:
    return verutils.Version(
        major=val['major'],
        minor=val['minor'],
        stage=getattr(verutils.VersionStage, val['stage'].upper()),
        stage_no=val['stage_no'],
        local=tuple(val['local']),
    )
